# coding: utf8

from multiprocessing import RLock


class QueryManager(object):

    lock = RLock()

    def __init__(self, queries_file: str, query_done_file: str):
        self.queries_file = queries_file
        self.query_done_file = query_done_file

    def load_query(self, data_file):
        queries = []
        with open(data_file, 'r', encoding="utf8") as f:
            for line in f:
                line = line.strip()
                if line == "" or line.startswith("#"):
                    continue
                queries.append(line)
        return queries

    def get_all_queries(self):
        self.lock.acquire()
        queries = self.load_query(self.queries_file)
        self.lock.release()
        return queries

    def get_query_done(self):
        self.lock.acquire()
        query_done = self.load_query(self.query_done_file)
        self.lock.release()
        return query_done

    def next_query(self):
        self.lock.acquire()
        queries = self.get_all_queries()
        query_done = set(self.get_query_done())
        self.lock.release()
        for query in queries:
            if query not in query_done:
                return query
        return None

    def query_done(self, query):
        self.lock.acquire()
        with open(self.query_done_file, 'a', encoding="utf8") as f:
            f.write("%s\n" % query)
        self.lock.release()



